class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution(object):
    def longestUnivaluePath(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        self.max_len = 0
        if not root:
            return 0

        def search(node):
            left_count = 0
            right_count = 0
            if node.left:
                left_val, left_count = search(node.left)
                if left_val == node.val:
                    left_count += 1
                else:
                    left_count = 0
            if node.right:
                right_val, right_count = search(node.right)
                if right_val == node.val:
                    right_count += 1
                else:
                    right_count = 0
            self.max_len = max(self.max_len, left_count + right_count)
            return node.val, max(left_count, right_count)

        search(root)
        return self.max_len